Add stream context #5
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tqchen
reviewed
Sep 14, 2025
python/tvm_ffi/stream.py
Outdated
| Examples | ||
| -------- | ||
| .. code-block:: python | ||
| s = torch.cuda.Stream() |
Member
There was a problem hiding this comment.
need indent one level and a blank line between code block and code
python/tvm_ffi/cython/base.pxi
Outdated
| return <uint64_t>prev_stream | ||
|
|
||
|
|
||
| class StreamContext: |
Member
There was a problem hiding this comment.
given this is not cython dependent, move out to stream.py
tqchen
reviewed
Sep 14, 2025
python/tvm_ffi/cython/base.pxi
Outdated
| TVMFFIStreamHandle stream, | ||
| TVMFFIStreamHandle* opt_out_original_stream) nogil | ||
|
|
||
| cdef _env_set_current_stream(int device_type, int device_id, uint64_t stream): |
Member
There was a problem hiding this comment.
we can expose this as def, so it can be called from python side
tqchen
reviewed
Sep 14, 2025
python/tvm_ffi/cython/base.pxi
Outdated
| DLTensor* TVMFFITensorGetDLTensorPtr(TVMFFIObjectHandle obj) nogil | ||
| DLDevice TVMFFIDLDeviceFromIntPair(int32_t device_type, int32_t device_id) nogil | ||
|
|
||
|
|
tqchen
approved these changes
Sep 15, 2025
tqchen
pushed a commit
that referenced
this pull request
Sep 25, 2025
This PR addresses a memory corruption issue in Cython. The fix involves
ensuring that the `ByteArrayArg` object, which holds the type key, is
properly destructed after being passed to the `TVMFFITypeKeyToIndex`
function. This prevents a potential read-after-free scenario, as
reported by ASan.
## ASan Report
```
READ of size 9 at 0x604000420a30 thread T0
...
#5 0x7fdb57299506 in __pyx_f_4core__type_info_create_from_type_key /home/dolores/Projects/tvm-ffi/build/core.cpp:17732
...
0x604000420a30 is located 32 bytes inside of 42-byte region [0x604000420a10,0x604000420a3a)
freed by thread T0 here:
...
#4 0x7fdb572994e2 in __pyx_f_4core__type_info_create_from_type_key /home/dolores/Projects/tvm-ffi/build/core.cpp:17731
...
previously allocated by thread T0 here:
...
#8 0x7fdb57299366 in __pyx_f_4core__type_info_create_from_type_key /home/dolores/Projects/tvm-ffi/build/core.cpp:17718
```
<img width="1444" height="904" alt="image"
src="https://github.com/user-attachments/assets/7a80d33d-dedf-41ca-ac77-108e63b8e57b"
/>
## Recommended ASan Options
One will need to preload `libasan` to properly work with CPython, and
`libstdc++` to properly intercept `__cxa_throw`. The path to those two
files can be found using:
```
ASAN="$(gcc -print-file-name=libasan.so)"
STDCXX="$(g++ -print-file-name=libstdc++.so.6)"
LD_PRELOAD="$ASAN $STDCXX"
```
Additionally, it might be helpful to tweak
```
PYTHONMALLOC=malloc
```
and run with ASan options
```
ASAN_OPTIONS="detect_leaks=0:abort_on_error=1:symbolize=1:fast_unwind_on_malloc=0"
```
Notably, turning on `detect_leaks=1` will lead to bunch of irrelevant
noisy reports. Better turning it off.
Merged
tqchen
pushed a commit
that referenced
this pull request
Oct 9, 2025
This PR exposes the get env stream method as followup of #5.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds the stream context into ffi, so that ffi env stream can be updated. The
tvm_ffi.use_torch_streamis for wrapping the torch stream/graph context. And lower-leveltvm_ffi.use_raw_streamis for creating context with existing stream handle.Example for
tvm_ffi.use_torch_stream:case with torch stream:
case with torch cuda graph
case with current stream by default
Eaxmple for
tvm_ffi.use_raw_stream: